/*
* This file is a part of the open-eBackup project.
* This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
* If a copy of the MPL was not distributed with this file, You can obtain one at
* http://mozilla.org/MPL/2.0/.
*
* Copyright (c) [2024] Huawei Technologies Co.,Ltd.
*
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
* EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
* MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
*/
#include <algorithm>
#include "PrefixUtils.h"
#include "NodeTaskManager.h"

using namespace std;
using namespace OBSPlugin;

NodePrefixEntity::NodePrefixEntity(const vector<PrefixInfo>& prefixInfo)
{
    for (auto prefix : prefixInfo) {
        ExpandPrefixInfo(prefix);
    }
}
void NodePrefixEntity::AddPrefixInfo(const PrefixInfo& prefixInfo)
{
    m_contents.insert(prefixInfo);
}

size_t NodePrefixEntity::EntityCount()
{
    return m_contents.size();
}

void NodePrefixEntity::Clear()
{
    m_contents.clear();
}

NodePrefixEntity NodePrefixEntity::operator - (const NodePrefixEntity & nodePrefixEntity) const
{
    NodePrefixEntity diff;
    std::set_difference(m_contents.begin(), m_contents.end(),
                        nodePrefixEntity.m_contents.begin(), nodePrefixEntity.m_contents.end(),
                        std::inserter(diff.m_contents, diff.m_contents.begin()),
                        NodePrefixEntity::Cmp);
    return diff;
}

NodePrefixEntity NodePrefixEntity::operator + (const NodePrefixEntity & nodePrefixEntity) const
{
    NodePrefixEntity merged;
    std::merge(m_contents.begin(), m_contents.end(),
               nodePrefixEntity.m_contents.begin(), nodePrefixEntity.m_contents.end(),
               std::inserter(merged.m_contents, merged.m_contents.begin()),
               NodePrefixEntity::Cmp);
    return merged;
}

void NodePrefixEntity::Remove(const NodePrefixEntity& nodePrefixEntity)
{
    for (auto iter = m_contents.begin(); iter != m_contents.end();) {
        auto item = std::find_if(nodePrefixEntity.m_contents.begin(), nodePrefixEntity.m_contents.end(),
            [&](const PrefixInfo &prefixInfo) {
                return prefixInfo.buildKey() == (*iter).buildKey();
            });
        if (item != nodePrefixEntity.m_contents.end()) {
            iter = m_contents.erase(iter);
        } else {
            iter++;
        }
    }
}

void NodePrefixEntity::AddPrefixInfoList(const std::vector<PrefixInfo>& prefixInfoList)
{
    m_contents.insert(prefixInfoList.begin(), prefixInfoList.end());
}

vector<PrefixInfo> NodePrefixEntity::GetPrefixInfo()
{
    std::vector<PrefixInfo> prefixInfoList;
    prefixInfoList.reserve(m_contents.size());
    prefixInfoList.assign(m_contents.begin(), m_contents.end());
    return prefixInfoList;
}

void NodePrefixEntity::ExpandPrefixInfo(const PrefixInfo& prefixInfo)
{
    for (auto parentPrefix : prefixInfo.parentPrefixFilter) {
        if (prefixInfo.subPrefixs.empty()) {
            PrefixInfo tempInfo(prefixInfo.bucketName, parentPrefix, "",
                                prefixInfo.isPrefixLevelObject);
            AddPrefixInfo(tempInfo);
            continue;
        }
        for (auto subPrefix : prefixInfo.subPrefixs) {
            PrefixInfo tempInfo(prefixInfo.bucketName, parentPrefix, subPrefix,
                                prefixInfo.isPrefixLevelObject);
            AddPrefixInfo(tempInfo);
        }
    }
}

PrefixEntityIterator::PrefixEntityIterator(NodePrefixEntity* entity)
    : m_nodePrefixEntity(entity), m_prefixIt(entity->m_contents.begin()) {}

void PrefixEntityIterator::Next()
{
    m_prefixIt++;
}

bool PrefixEntityIterator::End() const
{
    return m_prefixIt == m_nodePrefixEntity->m_contents.end();
}

bool PrefixEntityIterator::Empty() const
{
    return m_nodePrefixEntity->m_contents.empty();
}

void PrefixEntityIterator::operator ++ (int)
{
    Next();
}

PrefixInfo& PrefixEntityIterator::Current()
{
    return const_cast<PrefixInfo &>(*m_prefixIt);
}

vector<BackupObjectInfo> NodeTaskManager::AllocateNodePrefix(
    const vector<PrefixInfo>& curBucketPrefixs,
    const vector<string>& repoEsnList,
    vector<BackupObjectInfo>& lastObjectInfos)
{
    if (repoEsnList.size() == 0) {
        return vector<BackupObjectInfo>();
    }
    if (lastObjectInfos.empty()) { // 首次/全量备份首次分配前缀
        return FirstAllocateNode(curBucketPrefixs, repoEsnList);
    }
    // step 1: 合并上次副本前缀信息
    NodePrefixEntity curNodePrefixEntity(curBucketPrefixs); // 展开后当前扫描的桶前缀信息
    NodePrefixEntity prevNodePrefixEntity = MergeNodePrefix(lastObjectInfos); // 上次副本记录的桶前缀
    NodePrefixEntity newNodePrefixEntity = curNodePrefixEntity - prevNodePrefixEntity; // 本次新增的前缀

    // step 2: 先将新增的前缀平均分配到各个节点,该值为本次任务拆分结果。
    ReallocatePrefixNode(newNodePrefixEntity, lastObjectInfos);
    vector<BackupObjectInfo> newObjectInfoList = lastObjectInfos;
    
    // step 3: 删除已不存在的前缀信息,通过出差返回保存到Cache仓中
    NodePrefixEntity deletedNodePrefixEntity = prevNodePrefixEntity - curNodePrefixEntity; // 本次删除的前缀
    RemovePrefixNode(deletedNodePrefixEntity, lastObjectInfos);
    return newObjectInfoList;
}

void NodeTaskManager::RemovePrefixNode(NodePrefixEntity& toRemovePrefix,
    vector<BackupObjectInfo>& mergeObjectInfos)
{
    if (toRemovePrefix.EntityCount() == 0) {
        return;
    }
    for (auto iter = mergeObjectInfos.begin(); iter != mergeObjectInfos.end();) {
        auto &prefixInfoList = (*iter).prefixInfo;
        // flat PrefixInfo
        NodePrefixEntity flatPrefixInfo(prefixInfoList);
        flatPrefixInfo.Remove(toRemovePrefix);
        auto merged = MergeCollection(flatPrefixInfo.GetPrefixInfo());
        if (merged.size() != 0) {
            (*iter).prefixInfo.swap(merged);
            iter++;
        } else {
            iter = mergeObjectInfos.erase(iter);
        }
    }
}

void NodeTaskManager::ReallocatePrefixNode(NodePrefixEntity& toAllocatePrefix,
    vector<BackupObjectInfo>& mergeObjectInfos)
{
    if (toAllocatePrefix.EntityCount() == 0) {
        return;
    }
    uint32_t prefixCount = toAllocatePrefix.EntityCount();
    uint32_t esnCount = mergeObjectInfos.size();
    std::vector<uint32_t> nodePrefixAssign = AverageAssign(prefixCount, esnCount);
    uint32_t actualNodeNum = nodePrefixAssign.size();
    PrefixEntityIterator currIt(&toAllocatePrefix);
    NodePrefixEntity allocatePrefix;
    uint32_t alloctedNum = 0;
    uint32_t i = 0;
    while (!currIt.End() && i < actualNodeNum) {
        alloctedNum++;
        allocatePrefix.AddPrefixInfo(currIt.Current());
        if (alloctedNum % nodePrefixAssign[i] == 0) {
            allocatePrefix.AddPrefixInfoList(mergeObjectInfos[i].prefixInfo);
            auto merged = MergeCollection(allocatePrefix.GetPrefixInfo());
            mergeObjectInfos[i].prefixInfo.swap(merged);
            allocatePrefix.Clear();
            alloctedNum = 0;
            i++;
        }
        currIt++;
    }
}

std::vector<BackupObjectInfo> NodeTaskManager::FirstAllocateNode(const std::vector<PrefixInfo>& curBucketPrefixs,
    const std::vector<std::string>& repoEsnList)
{
    vector<BackupObjectInfo> resultObjectInfoList;
    vector<string> esnList = repoEsnList;
    uint32_t esnCount = esnList.size();
    sort(esnList.begin(), esnList.end());
    NodePrefixEntity curNodePrefixEntity(curBucketPrefixs); // 展开后的桶前缀信息
    uint32_t prefixCount = curNodePrefixEntity.EntityCount();
    std::vector<uint32_t> nodePrefixAssign = AverageAssign(prefixCount, esnCount);
    uint32_t actualNodeNum = nodePrefixAssign.size();
    for (uint32_t i = 0; i < actualNodeNum; i++) {
        BackupObjectInfo backupObjectInfo;
        backupObjectInfo.esn = esnList[i];
        resultObjectInfoList.emplace_back(backupObjectInfo);
    }
    // 平均分配桶前缀到对应节点ESN，多余的前缀统一放最后一个
    ReallocatePrefixNode(curNodePrefixEntity, resultObjectInfoList);
    return resultObjectInfoList;
}

vector<uint32_t> NodeTaskManager::AverageAssign(uint32_t prefixCount, uint32_t nodeCount)
{
    if (nodeCount == 0) {
        return vector<uint32_t>();
    }
    uint32_t actualNodeNum = nodeCount > prefixCount ? prefixCount : nodeCount;
    uint32_t everyTaskPrefixNum = prefixCount / nodeCount;
    uint32_t remainNum = prefixCount % nodeCount;
    std::vector<uint32_t> newAllocatePrefixNum(actualNodeNum, everyTaskPrefixNum);
    if (actualNodeNum >= 1) {
        newAllocatePrefixNum[actualNodeNum - 1] += remainNum;
    }
    return newAllocatePrefixNum;
}

NodePrefixEntity NodeTaskManager::MergeNodePrefix(const std::vector<BackupObjectInfo>& objectInfos)
{
    NodePrefixEntity nodePrefixEntity;
    for (auto backupObjectInfo : objectInfos) {
        nodePrefixEntity = nodePrefixEntity + NodePrefixEntity(backupObjectInfo.prefixInfo);
    }
    return nodePrefixEntity;
}
